import dashscope
from dashscope import ImageSynthesis
from http import HTTPStatus
import os
import requests
import json
from urllib.parse import urlparse, unquote
from pathlib import PurePosixPath

class Qwen:
    def __init__(self, api_key):
        self.api_key = api_key
        dashscope.api_key = api_key
    
    def generate_image(self, prompt, style, sketch_path):
        # 生成图片
        rsp = ImageSynthesis.call(
            model="wanx-sketch-to-image-lite",
            prompt=prompt,
            style=style,
            sketch_image_url=sketch_path,
            task="image2image",
            n=1,
            size="768*768"
        )
        if rsp.status_code == HTTPStatus.OK:
            return {"status": "success", "url": rsp.output.results[0]["url"],"task_id": rsp.output.task_id}
        else:
            return {"status": "error", "message": rsp.message}
        
    def recognize_image(self, prompt, image_path):
        # 识别图片
        messages = [
        {
            "role": "system",
            "content": [
            {"text": "You are a helpful assistant."}]
        },
        {
            "role": "user",
            "content": [
            {"image": image_path},
            {"text": prompt}]
        }]
        response = dashscope.MultiModalConversation.call(
            api_key = os.getenv(self.api_key),
            model = 'qwen-vl-max-latest',
            messages = messages
        )
        return response.output.choices[0].message.content[0]["text"]
    
    def extend_image(self, prompt, image_path):
        # 扩展画面
        url = 'https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/out-painting'
        headers = {
            'X-DashScope-Async': 'enable',
            'Authorization': 'Bearer ' + self.api_key,
            'Content-Type': 'application/json'
        }
        data = {
            "model": "image-out-painting",
            "input": {
                "image_url": image_path
            },
            "parameters": {
                "x_scale": 1.5,
                "y_scale": 1.5
            }
        }
        response = requests.post(url, headers=headers, data=json.dumps(data))

        task_id = response.output.task_id
        url = f'https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}'
        headers = {
            'Authorization': 'Bearer ' + self.api_key 
        }

        response = requests.get(url, headers=headers)
        return response.output.outpu_image_url
    
    def redraw_image(self, prompt, base_image_path, mask_image_path):
        # 局部重绘
        model = "wanx-x-painting"
        task = "image2image"
        extra_input = {
            "base_image_url": base_image_path,
            "mask_image_url": mask_image_path
        }


        print('----sync call, please wait a moment----')
        rsp = ImageSynthesis.call(model=model,
                                prompt=prompt,
                                n=1,
                                size='1024*1024',
                                task=task,
                                extra_input=extra_input)
        if rsp.status_code == HTTPStatus.OK:
            return rsp.output
        else:
            print('sync_call Failed, status_code: %s, code: %s, message: %s' %
                (rsp.status_code, rsp.code, rsp.message))